########################
#
# Copyright 2010, Vanderbilt University
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# 
#       http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an "AS IS"
# BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
# or implied.  See the License for the specific languaged governing 
# permissions and limitations under the License.
#
########################

from bitarray import bitarray
from itertools import izip, chain
import util
import settings

defaultBitweights = sum([[(len(field.possibleValues.keys())-1)]*(len(field.possibleValues.keys())-1) for field in settings.fields], [])
alpha = 1
beta = 1

# BASIC FUNCTIONS
def policySize(policy):
    return policy.length()

def mostSpecificPolicy(len):
    """
    All ones

    """
    return len * bitarray('1')

def mostGeneralPolicy(len):
    """
    All zeroes 

    """
    return len * bitarray('0')

def randomPolicy(len):
    """
    Creates a random bitstring of the appropriate length.  This will
    be used in testing
    
    """
    #import random
    #return random.getrandbits(len)
    return bitarray(len) 

   
def safeHarbor():
    policy = bitarray()
    for field in settings.fields:
        if field.label == "BIRTH_YEAR":
            bools = [field.possibleValues[key][0] <= 90 \
                     for key \
                     in sorted(field.possibleValues.keys())]
            policy.extend(bitarray(bools[:-1]))
        elif field.label == "AGE":
            bools = [field.possibleValues[key][0] < 90 \
                     for key \
                     in sorted(field.possibleValues.keys())]
            policy.extend(bitarray(bools[:-1]))
        elif field.label == "ZIP":
            bools = []
            prevZip = False
            for zip in sorted(field.possibleValues.keys()):
               if prevZip and prevZip[:3] == zip[:3]:
                    bools.append(0)
                    prevZip = zip
               elif prevZip and prevZip[:3] != zip[:3]:
                    bools.append(1)
                    prevZip = zip
               else:
                    prevZip = zip
            policy.extend(bitarray(bools))
        else:  
            policy.extend((len(field.possibleValues.keys())-1) * bitarray('1'))
    return policy
 
def otherPresetPolicies(generalizedRace, ageIncrement, ageCutoff, zipLevel):
    """
    Written to enable research done with pre-defined policies to be updated
    using the new framework
    
    Previous policies investigated included:
    Gender, Race, Age- Full specificity
    Gender, Race, Age in 5 year increments
    Gender, Race, Age in 10 year increments
    Gender, Generalized Race, Age
    Gender, Generalized Race, Age in 5 Year increments
    Gender, Generalized Race, Age in 10 Year increments

    Essentially, this function will enable me to hard-code the standard
    generalized race (black/white/other) and allow for any regular division
    of ages.
    """
    policy = bitarray()
    ageFn = util.ageLookup(ageIncrement, ageCutoff)   
    for field in settings.fields:
        if field.label == "BIRTH_YEAR" or field.label == "AGE":
            bools = []
            prevAgeRange = False
            for age in sorted(field.possibleValues.keys()):
                if prevAgeRange and prevAgeRange == ageFn(int(age))[:2]:
                    bools.append(0)
                elif prevAgeRange and prevAgeRange != ageFn(int(age))[:2]:
                    bools.append(1)
                prevAgeRange = ageFn(int(age))[:2]
            policy.extend(bitarray(bools))
        elif field.label == "RACE" and generalizedRace:
            bools = []
            prevRace = False
            for race in sorted(field.possibleValues.keys()):
                if prevRace and \
                ("White alone" in field.possibleValues[race] or \
                "Black or African American alone" in field.possibleValues[race] or \
                "White alone" in field.possibleValues[prevRace] or \
                "Black or African American alone" in field.possibleValues[prevRace]):
                    bools.append(1)
                elif prevRace:
                    bools.append(0)
                prevRace = race
            policy.extend(bitarray(bools))
        elif field.label == "ZIP" and zipLevel < 5:
            bools = []
            prevZip = False
            for zip in sorted(field.possibleValues.keys()):
               if prevZip and prevZip[:zipLevel] == zip[:zipLevel]:
                    bools.append(0)
                    prevZip = zip
               elif prevZip and prevZip[:zipLevel] != zip[:zipLevel]:
                    bools.append(1)
                    prevZip = zip
               else:
                    prevZip = zip
            policy.extend(bitarray(bools))
        else:  
            policy.extend((len(field.possibleValues.keys())-1) * bitarray('1'))
    return policy
 
# Policy Relationships
def isAncestor(ancestor, descendant):
    """
    Ancestor is an ancestor of descendant  iff all the pplaces where descendant has 
    a zero, ancestor also has a zero

    """
    if ancestor == descendant:
        return False
    return (~ancestor | descendant).all()     

def policyNeighbors(initialPolicy, more=True, less=True):
    """
    Generates a list of policies which are neighbors to the
    original policy, that is, those that differ in one place.
    This function is used as the parent for both generalizations
    (less=True, more=False), and specifications (more=True, 
    less=False) 
    
    """
    if not more and not less:
        return []
    alternatives = []
    for position in range(len(initialPolicy)):
        if less and initialPolicy[position]:
            alternative = bitarray(initialPolicy)
            alternative[position] = 0
            alternatives.append(alternative)
        if more and not initialPolicy[position]:
            alternative = bitarray(initialPolicy)
            alternative[position] = 1
            alternatives.append(alternative)
    return alternatives 

def policyGeneralizations(initialPolicy):
    """
    Generates a list of policies which are one level generalizations
    of the initial policy (i.e. have one position which was one in the
    original and is zero now
    
    """
    return policyNeighbors(initialPolicy, less=True, more=False)

def policySpecifications(initialPolicy):
    """
    Generates a list of policies which are one level specifications 
    of the initial policy (i.e. have one position which was zero in the
    original and is one now
    
    """
    return policyNeighbors(initialPolicy, less=False, more=True)

def randomChild(initialPolicy):
    import random
    try:
        child =  random.choice(policySpecifications(initialPolicy))
        return child
    except IndexError:
        print "Policy has no children:", initialPolicy
        raise

def randomParent(initialPolicy):
    import random
    try:
        parent = random.choice(policyGeneralizations(initialPolicy))
        return parent 
    except IndexError:
        print "Policy has no parents:", initialPolicy
        raise
    
def shortcutRepr(policy):
    """
    Represents all values that are grouped together.
    
    """
    position = 0
    groups = []
    for field in settings.fields:
        span = len(field.possibleValues.keys()) - 1
        divisions = policy[position:position+span]
        position += span
        generalizations = divisions.search([False]) 
        values = sorted(field.possibleValues.keys())
        prevGen = -7 
        currentList = [] 
        for gen in generalizations:
            if gen == prevGen + 1:
                currentList.append(values[gen+1])
            else:
                if len(currentList) > 0:
                    groups.append(currentList)
                currentList = [values[gen], values[gen+1]]
            prevGen = gen
        if len(currentList) > 0:
            groups.append(currentList)
    return groups    
            
def shortcutToPolicy(listOfLists):
    policy = bitarray()
    for field, groups in zip(settings.fields, listOfLists):
        fieldPolicy = bitarray('1') * (len(field.possibleValues.keys()) - 1)
        values = sorted(field.possibleValues.keys())
        for group in groups:
            minPos = values.index(group[0])
            maxPos = values.index(group[-1])
            fieldPolicy[minPos:maxPos] = (maxPos - minPos) * bitarray([False])
        policy.extend(fieldPolicy)
    return policy

def policyToString(policy):
    position = 0
    string = ""
    for field in settings.fields:
        span = len(field.possibleValues.keys()) - 1
        divisions = policy[position:position+span]
        position += span 
        q = chain(*izip(sorted(field.possibleValues.keys()), \
                       [" | " if d else ", " for d in divisions]))
        string += "["+"".join(q)+sorted(field.possibleValues.keys())[-1]+"]"
    return string    
    #position = 1
    #division = ""
    #for field in settings.fields: 
    #     position -= 1
    #     division += "["
    #     for value in sorted(field.possibleValues.keys()):
    #        division += str(value)
    #        if bitstring & (1<<position):
    #        #if bitstring & (1>>position):
    #            division += "]["
    #        else:
    #            division += ", "
    #        position += 1
    #     division = division[:-2] 
    #     division += "]"
    #return division

def policyValidation(policy):
    """
    Validates random policies generated by the exploration process,
    if valid, the original policy will be returned.  If invalid, the
    "closest" valid policy (for each field) will be returned.
    
    """
    pass
    
def policyLevelsBetween(min, max):
    return (min ^ max).count()
def policyHalfpointOrig(min, max):
    """
    Generates a random policy halfway between min and max.  Change
    half the zeros in min that are ones in max to ones.
    
    """
    import random
    diff = (~min) & max
    places = diff.count(True)
    positions = random.sample(diff.search([True]), places / 2)
    for p in positions:
        diff[p] = 0 
    return min | diff 

def policyHalfpoint(max, min):
    """
    Generates a random policy halfway between min and max.  Change
    ones to zeros.  This works with our weighting system. 
    """
    import random
    diff = (~max) & min
    places = diff.count(True)
    mask = mostSpecificPolicy(policySize(min))
    positions = random.sample(diff.search([True]), places / 2)
    for p in positions:
        mask[p] = 0 
    return min & mask 


def directedPolicyHalfpoint(max, min, bitweights=defaultBitweights):
    """
    Generates a random policy halfway between min and max.  Change
    half the zeros in min that are ones in max to ones.
    
    """
    diff = (~max) & min
    places = diff.search([True])
    mask = mostSpecificPolicy(policySize(min))
    positions = util.weightedSample(zip([bitweights[p] for p in places], places), 
                                    len(places) / 2)
    for p in positions:
        mask[p] = 0
    return min & mask 

def sizeOfFieldWeights():
    return sum([[(len(field.possibleValues.keys())-1)]*(len(field.possibleValues.keys())-1) for field in settings.fields], [])

def prevSolnWeights(prevPoints):
    n = len(prevPoints[0])
    return [sum([e[k] for e in prevPoints]) + 1 for k in range(n)]

def informedMidpoint(len, prevPoints):
    """
    Create a set of weights for informed selection.  These weights include 
    information about the existing solutions and (possibly) the size of the domains
    of each attribute
    """
    notDividedYet = reduce(or_, [~policy for policy in prevPoints])
    if notDividedYet.search([True]) < len/2:
        #No possible solution exists
        return None
    positions = random.sample(notDividedYet.search([True]), len/2)
    initial = bitarray([1]*len)
    for p in positions:
        initial[p] = 0
    return initial 

def directedWithPrevKnowledgeHalfpoint(min, max, prevPoints):
    """
    Count the number of times something has been divided.  If it has been
    divided previously, then we should choose it with a lower frequency
    (lower weight).
    
    """
    prevCount = [sum([not ele[k] for ele in prevPoints]) for k in range(len(prevPoints[0]))] 
    weights = [(1/(count+1.))*len(prevPoints) for count in prevCount]
    integratedWeights = [w1**alpha + w2**beta for (w1, w2) in zip(weights, defaultBitweights)]
    return directedPolicyHalfpoint(min, max, integratedWeights)

def weightedHalfpoint(min, max, prevPoints):
    weights = [w1**alpha + w2 ** beta for (w1, w2) in zip(sizeOfFieldWeights(), prevSolnWeights(prevPoints))]
    return directedPolicyHalfpoint(min, max, weights)

